import struct
from collections import namedtuple
import re
import sys

# 定义 PostgreSQL 消息类型
PGMessageType = {
    b"R": "Authentication",
    b"K": "BackendKeyData",
    b"2": "BindComplete",
    b"3": "CloseComplete",
    b"C": "CommandComplete",
    b"d": "CopyData",
    b"c": "CopyDone",
    b"G": "CopyInResponse",
    b"H": "CopyOutResponse",
    b"W": "CopyBothResponse",
    b"D": "DataRow",
    b"I": "EmptyQueryResponse",
    b"E": "ErrorResponse",
    b"V": "FunctionCallResponse",
    b"n": "NoData",
    b"N": "NoticeResponse",
    b"A": "NotificationResponse",
    b"t": "ParameterDescription",
    b"S": "ParameterStatus",
    b"1": "ParseComplete",
    b"s": "PortalSuspended",
    b"Z": "ReadyForQuery",
    b"T": "RowDescription",
    b"P": "Parse",
    b"B": "Bind",
    b"D": "Describe",
    b"Q": "Query",
}

# 定义列描述结构
ColumnDescriptor = namedtuple(
    "ColumnDescriptor",
    [
        "name",
        "table_oid",
        "column_attrnum",
        "type_oid",
        "type_size",
        "type_modifier",
        "format_code",
    ],
)

# 添加常见PostgreSQL类型OID映射
PG_TYPE_OID = {
    16: "bool",
    20: "int8",
    21: "int2",
    23: "int4",
    25: "text",
    700: "float4",
    701: "float8",
    1043: "varchar",
    1082: "date",
    1114: "timestamp",
    1184: "timestamptz",
    1700: "numeric",
}

"""
类型长度解释:
+--------+------------+------------+------------------------------+
| OID    | 类型名     | type_size  | 解释                         |
+--------+------------+------------+------------------------------+
| 23     | int4       | 4          | 4 字节固定长度               |
| 1043   | varchar    | -1         | 可变长度字符串               |
| 1700   | numeric    | -1         | 精度可变的数值类型           |
| 17     | bytea      | -1         | 可变长度二进制               |
| 1082   | date       | 4          | 固定 4 字节                  |
| 1083   | time       | 8          | 固定长度时间（含微秒）       |
| 1114   | timestamp  | 8          | 固定长度时间戳               |
| 25     | text       | -1         | 可变长度文本                 |
| 21     | int2       | 2          | 2 字节                       |
| 20     | int8       | 8          | 8 字节                       |
| 16     | bool       | 1          | 1 字节布尔值                 |
+--------+------------+------------+------------------------------+
"""


def parse_data_row(data, pos=0):
    """解析数据行消息 (类型 'D')"""
    if len(data) - pos < 2:
        return [], pos

    # 读取列数
    num_columns = struct.unpack("!H", data[pos : pos + 2])[0]
    pos += 2

    row = []
    for _ in range(num_columns):
        if len(data) - pos < 4:
            break

        # 读取列长度
        col_len = struct.unpack("!I", data[pos : pos + 4])[0]
        pos += 4

        if col_len == -1:  # NULL 值
            row.append(None)
            continue

        if len(data) - pos < col_len:
            break

        # 读取列值 (作为文本)
        value = data[pos : pos + col_len].decode("utf-8")
        pos += col_len
        row.append(value)

    return row, pos


def parse_pg_message(data, pos=0):
    """解析 PostgreSQL 消息"""
    messages = []

    while pos < len(data):
        if len(data) - pos < 5:
            break

        msg_type = data[pos : pos + 1]
        # 处理消息长度(包含自身4字节)
        msg_len = struct.unpack("!I", data[pos + 1 : pos + 5])[0]
        pos += 5

        if pos + msg_len - 4 > len(data):  # msg_len包含自身4字节
            break

        msg_data = data[pos : pos + msg_len - 4]
        pos += msg_len - 4

        # 根据消息类型解析内容
        if msg_type == b"T":  # 行描述
            columns, _ = parse_row_description(msg_data)
            messages.append(("RowDescription", columns))
        elif msg_type == b"D" and len(msg_data) > 2:  # 数据行
            row, _ = parse_data_row(msg_data)
            messages.append(("DataRow", row))
        elif msg_type == b"C":  # 命令完成
            tag = msg_data.decode("utf-8")
            messages.append(("CommandComplete", tag))
        elif msg_type == b"Z":  # 准备查询
            status = msg_data.decode("utf-8")
            messages.append(("ReadyForQuery", status))
        else:
            msg_name = PGMessageType.get(msg_type, f"Unknown({msg_type.hex()})")
            messages.append((msg_name, msg_data.hex()))

    return messages, pos


def parse_row_description(data, pos=0):
    """解析行描述消息 (类型 'T')"""
    if len(data) - pos < 2:
        return [], pos

    # 读取列数
    num_columns = struct.unpack("!H", data[pos : pos + 2])[0]
    pos += 2

    columns = []
    for _ in range(num_columns):
        # 检查剩余数据是否足够
        if len(data) - pos < 19:  # 列名至少1字节+18字节属性
            break

        # 读取列名 (以null结尾)
        try:
            name_end = data.index(b"\x00", pos)
        except ValueError:
            break

        name = data[pos:name_end].decode("utf-8")
        pos = name_end + 1

        # 检查剩余数据是否足够18字节
        if len(data) - pos < 18:
            print(f"警告: 需要18字节但只有{len(data)-pos}字节可用于列属性")
            break

        # 读取列属性
        try:
            (
                table_oid,
                column_attrnum,
                type_oid,
                type_size,
                type_modifier,
                format_code,
            ) = struct.unpack("!IhIhih", data[pos : pos + 18])
            pos += 18
        except struct.error as e:
            print(f"解析列属性错误: {e}, 位置: {pos}, 剩余数据: {len(data)-pos}字节")
            break

        columns.append(
            ColumnDescriptor(
                name,
                table_oid,
                column_attrnum,
                type_oid,
                type_size,
                type_modifier,
                format_code,
            )
        )

    return columns, pos


def read_data_from_file(filename):
    """从文件读取数据，支持二进制或十六进制"""
    try:
        with open(filename, "rb") as f:
            content = f.read()
            # 尝试检测是否为十六进制文本
            if b" " in content or b"\n" in content:
                hex_str = re.sub(rb"[^0-9a-fA-F]", b"", content)
                return bytes.fromhex(hex_str.decode("ascii"))
            return content
    except FileNotFoundError:
        print(f"错误: 文件 '{filename}' 未找到")
        sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("用法: python3 pg_parser.py <data_file>")
        sys.exit(1)

    filename = sys.argv[1]
    binary_data = read_data_from_file(filename)

    print(f"正在解析数据，总长度: {len(binary_data)}字节")
    messages, _ = parse_pg_message(binary_data)

    # 更友好的输出格式
    for msg_type, content in messages:
        if msg_type == "RowDescription":
            print("\n=== 表结构 ===")
            for col in content:
                type_name = PG_TYPE_OID.get(col.type_oid, str(col.type_oid))
                print(
                    f"{col.name:15} {type_name:10} (OID:{col.type_oid}, 长度:{col.type_size})"
                )
        elif msg_type == "DataRow":
            print("\n=== 行数据 ===")
            for i, val in enumerate(content, 1):
                print(f"列{i}: {val if val is not None else 'NULL'}")
        else:
            print(
                f"\n[{msg_type}] {content if isinstance(content, str) else content.hex()}"
            )
